import torch.nn as nn

class FoodCNN(nn.Module):
    def __init__(self, num_classes=11):
        super(FoodCNN, self).__init__()
        # 特征提取网络
        self.cnn_layers = nn.Sequential(
            # 第一卷积块 [3, 128, 128] -> [64, 128, 128]
            nn.Conv2d(3, 4, 3, padding=1),
            nn.BatchNorm2d(4),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # [64, 64, 64]

            # 第二卷积块 [64, 64, 64] -> [128, 64, 64]
            nn.Conv2d(4, 8, 3, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # [128, 32, 32]

            # 第三卷积块 [128, 32, 32] -> [256, 32, 32]
            nn.Conv2d(8, 16, 3, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  # [256, 16, 16]

            # 第四卷积块 [256, 16, 16] -> [512, 16, 16]
            nn.Conv2d(16, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d(4, 4),  # [512, 4, 4]
        )

        # 分类器
        self.fc_layers = nn.Sequential(
            nn.Linear(32 * 4 * 4, 64),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(32, num_classes)
        )

    def forward(self, x):
        x = self.cnn_layers(x)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc_layers(x)
        return x
